package utils

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
)

var AesUtils = aesUtils{
	BlockSize: 16, // AES 块大小固定为 16 字节
}

type aesUtils struct {
	BlockSize int
}

// GenerateIV 生成 16 字节的随机 IV，并以 16 进制字符串返回
func (u aesUtils) GenerateIV() (string, error) {
	iv := make([]byte, u.BlockSize)
	_, err := io.ReadFull(rand.Reader, iv)
	if err != nil {
		return "", fmt.Errorf("生成 IV 失败: %v", err)
	}
	return hex.EncodeToString(iv), nil
}

// GenerateKey 生成 16 字节（128 位）的随机 AES 密钥，并以 16 进制字符串返回
func (u aesUtils) GenerateKey() (string, error) {
	key := make([]byte, u.BlockSize) // 默认 AES-128
	_, err := io.ReadFull(rand.Reader, key)
	if err != nil {
		return "", fmt.Errorf("生成密钥失败: %v", err)
	}
	return hex.EncodeToString(key), nil
}

// Generate 返回一个随机 AES 密钥和 IV，均为 16 进制字符串组成的 Pair (key, iv)
func (u aesUtils) Generate() (keyHex string, ivHex string, err error) {
	keyHex, err = u.GenerateKey()
	if err != nil {
		return "", "", err
	}
	ivHex, err = u.GenerateIV()
	if err != nil {
		return "", "", err
	}
	return keyHex, ivHex, nil
}

// Encrypt 使用 AES-CBC/PKCS7Padding 加密明文，返回 16 进制编码的密文
// keyHex 和 ivHex 均为 16 进制字符串
func (u aesUtils) Encrypt(plaintext string, keyHex string, ivHex string) (string, error) {
	key, err := hex.DecodeString(keyHex)
	if err != nil {
		return "", fmt.Errorf("无效的密钥 Hex 字符串: %v", err)
	}
	iv, err := hex.DecodeString(ivHex)
	if err != nil {
		return "", fmt.Errorf("无效的 IV Hex 字符串: %v", err)
	}

	if len(key) != 16 {
		return "", errors.New("仅支持 16 字节 (128 位) AES 密钥，请修改代码以支持 24/32 字节")
	}
	if len(iv) != u.BlockSize {
		return "", fmt.Errorf("IV 必须是 %d 字节", u.BlockSize)
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("创建 AES cipher 失败: %v", err)
	}

	// PKCS7 Padding
	plaintextBytes := []byte(plaintext)
	plaintextBytes = u.pkcs7Pad(plaintextBytes, block.BlockSize())

	mode := cipher.NewCBCEncrypter(block, iv)
	ciphertext := make([]byte, len(plaintextBytes))
	mode.CryptBlocks(ciphertext, plaintextBytes)

	// 返回 16 进制编码的密文
	return hex.EncodeToString(ciphertext), nil
}

// Decrypt 使用 AES-CBC/PKCS7Padding 解密 16 进制密文，返回原始明文字符串
func (u aesUtils) Decrypt(ciphertextHex string, keyHex string, ivHex string) (string, error) {
	key, err := hex.DecodeString(keyHex)
	if err != nil {
		return "", fmt.Errorf("无效的密钥 Hex 字符串: %v", err)
	}
	iv, err := hex.DecodeString(ivHex)
	if err != nil {
		return "", fmt.Errorf("无效的 IV Hex 字符串: %v", err)
	}
	ciphertext, err := hex.DecodeString(ciphertextHex)
	if err != nil {
		return "", fmt.Errorf("无效的密文 Hex 字符串: %v", err)
	}

	if len(key) != u.BlockSize {
		return "", errors.New("仅支持 16 字节 (128 位) AES 密钥，请修改代码以支持 24/32 字节")
	}
	if len(iv) != u.BlockSize {
		return "", fmt.Errorf("IV 必须是 %d 字节", u.BlockSize)
	}

	block, err := aes.NewCipher(key)
	if err != nil {
		return "", fmt.Errorf("创建 AES cipher 失败: %v", err)
	}

	if len(ciphertext)%block.BlockSize() != 0 {
		return "", errors.New("密文长度必须是块大小的倍数")
	}

	mode := cipher.NewCBCDecrypter(block, iv)
	plaintext := make([]byte, len(ciphertext))
	mode.CryptBlocks(plaintext, ciphertext)

	// 去除 PKCS7 Padding
	plaintext, err = u.pkcs7Unpad(plaintext)
	if err != nil {
		return "", fmt.Errorf("去除 PKCS7 Padding 失败: %v", err)
	}
	return string(plaintext), nil
}

// ======================
// PKCS7 Padding / Unpadding
// ======================

func (u aesUtils) pkcs7Pad(data []byte, blockSize int) []byte {
	padding := blockSize - len(data)%blockSize
	padText := u.bytesRepeat(byte(padding), padding)
	return append(data, padText...)
}

func (u aesUtils) pkcs7Unpad(data []byte) ([]byte, error) {
	if len(data) == 0 {
		return nil, errors.New("空数据无法去填充")
	}
	padding := int(data[len(data)-1])
	if padding == 0 || padding > len(data) {
		return nil, errors.New("无效的 PKCS7 填充")
	}
	// 检查最后 padding 个字节是否都等于 padding
	for i := len(data) - padding; i < len(data); i++ {
		if int(data[i]) != padding {
			return nil, errors.New("无效的 PKCS7 填充")
		}
	}
	return data[:len(data)-padding], nil
}

// 辅助函数：生成 n 个 b 字节
func (u aesUtils) bytesRepeat(b byte, n int) []byte {
	buf := make([]byte, n)
	for i := 0; i < n; i++ {
		buf[i] = b
	}
	return buf
}
